package com.zhxq.ai.repository;

import cn.hutool.core.collection.CollUtil;
import com.zhxq.ai.domain.ChatMessage;
import com.zhxq.ai.domain.vo.MessageVo;
import com.zhxq.common.utils.DateUtils;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.stereotype.Component;

import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

@Component
public class MongoChatMemory implements ChatMemory {

    @Autowired
    private MongoTemplate mongoTemplate;

    @Override
    public void add(String conversationId, List<Message> messages) {
        Query query = new Query();
        query.addCriteria(Criteria.where("chatId").is(conversationId));
        ChatMessage chatMessage = mongoTemplate.findOne(query, ChatMessage.class);
        // 遍历添加的message，转换为MessageVo
        List<MessageVo> messageVos = messages.stream().map(item -> {
            MessageVo messageVo = new MessageVo(item);
            messageVo.setChatTime(DateUtils.getNowDate());
            return messageVo;
        }).toList();
        if(chatMessage == null) {
            // 没有会话，则插入
            chatMessage = new ChatMessage();
            chatMessage.setChatId(conversationId);
            chatMessage.setMessageList(messageVos);
            chatMessage.setUpdateTime(DateUtils.getNowDate());
            mongoTemplate.insert(chatMessage);
        } else {
            // 有会话，则增加新消息
            chatMessage.getMessageList().addAll(messageVos);
            chatMessage.setUpdateTime(DateUtils.getNowDate());
            mongoTemplate.save(chatMessage);
        }
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        Query query = new Query();
        query.addCriteria(Criteria.where("chatId").is(conversationId));
        ChatMessage chatMessage = mongoTemplate.findOne(query, ChatMessage.class);
        if(chatMessage == null || CollUtil.isEmpty(chatMessage.getMessageList())) {
            return Collections.emptyList();
        }
        return chatMessage.getMessageList().stream().map(messageVo -> {
            return switch (messageVo.getRole()) {
                case "user" -> new UserMessage(messageVo.getContent());
                case "assistant" -> new AssistantMessage(messageVo.getContent());
                case "system" -> new SystemMessage(messageVo.getContent());
                default -> throw new IllegalArgumentException("未知的消息类型");
            };
        }).collect(Collectors.toList());
    }

    @Override
    public void clear(String conversationId) {
        Query query = new Query();
        query.addCriteria(Criteria.where("chatId").is(conversationId));
        mongoTemplate.remove(query, ChatMessage.class);
    }
}
